#Frequency Counter Object - Mark Baggett @markbaggett
#This object is used to count the frequency of characters appearance in text.
#It was written for use with th SQL Injection tool and distributed with that tool.

import os

class FreqCounter(dict):
  """This object is used for counting the frequency of characters that follow other characters. It supports saving the data structure and loading it from disk.  The printtable() method is used to generate python source code to lookup the characters in frequency order.  Other key methods include tallyfile(), tallydict(), printraw(), save(), load(), probability() and resetcounts().  Check the help on each of them for more details. 

Here is an example of how this object might be used.

>>> c1=FreqCounter()    - Create a new Frequency counter object named c1
>>> c1.load("sqlinjector-default.freq")      - Load the default frequency counts
>>> c1.resetcounts  - reset all the characters to a frequency of 1
>>> c1.tallyfile("/Users/markbaggett/textdocuments/mobydick.txt")   - Add the character count from Moby Dick
1191463
>>> c1.tallyfile("/Users/markbaggett/textdocuments/lilwomen.txt")  - Add more characters..
1042048
>>> c1.tallyfile("/Users/markbaggett/textdocuments/war+peace.txt",weight=10)  - count each character in war and peace as 10 characters
3202941
>>> c1.save("sqlinjector-mytable.freq")  - Save a new frequency table for later use.
>>> c2=FreqCounter()    -  Create a c2 instance of FreqCounter()
>>> c2.tallychars="abcdefghijklmnopqrstuvwxyz"    - specify the character sets that the "tally" functions will count.
>>> c2.tallydict(c1)    - Tally all the characters in c1 into c2 (reducing the table to just those things in tallychar)
>>> c2.save("sqlinjector-myletters.freq")
>>> c2.probability("q","u")    - returns 98% chance that u follows q
>>> c2.lookup("q")  - returns a string containing the characters in order of frequency
>>> c2.printtable()  - prints a python function for looking up the string returnded by lookup()
>>> c2.printraw()
"""

  tallychars = "abcdefghijklmnopqrstuvwxyz0123456789.,_-$"
  ignorecase = True

  def __getitem__(self, item):
    try:
      return dict.__getitem__(self, item)
    except KeyError:
      value = self[item] = type(self)()
    return value


  def tallyfile(self,filename,weight=1):
    """tallyfile() accepts two parameters.  A filename and optionally you can specify a weight. Tallyfile will read the supplied fully qualified file path and tally up the character frequencies that appear in that file ADDING the characters in that table to the current values in the frequency counter object.  Weight will multiply the number of times it sees a character by this value.  So if you set weight to 10 one instance of b following f will be counted 10 times.    

Examples:   
counter.tallyfile("/home/user/file.txt", weight=100) - read the character frequencies in file.txt from the users home directory.  Count each instance of a character 100 times when adding the values to the character frequency count.

counter.tallyfile("/home/user/file.txt") - read the character frequencies in file.txt from the users home directory.  Use the default weight and count each instance of a character 1 times when adding the values to the character frequency count.""" 

    wordcount=0
    file=open(filename, 'r')
    linesinfile=file.readlines()
    for line in linesinfile:
      if self.ignorecase:
        line=line.lower()
      wordcount=wordcount+len(line)
      for char in range(len(line)-1):
        if not line[char] in self.tallychars or not line[char+1] in self.tallychars:
           continue
        if line[char+1] in self[line[char]]:
          self[line[char]][line[char+1]]=self[line[char]][line[char+1]]+(1*weight)
        else:
          self[line[char]][line[char+1]]=weight
    file.close()
    return wordcount

  def tallydict(self,objtomerge,weight=1):
    """tallydict() accepts two parameters.  Another frequency counter object and optionally you can specify a weight. Tallyfile will add values in the provided object to the current values in the frequency counter object.  Weight will multiply the number of times it sees a character by this value.  So if you set weight to 10 one instance of b following f will be counted 10 times.  If the weight parameter is not provided a default value of 1 will be used.

Example:   
counter.tallydict(counter2, weight=100) - Add character frequencies in counter2 to counter.  Count each instance of a character 100 times when adding the values to the character frequency count.

counter.tallydict(counter2) - Add character frequencies in counter2 to counter. 
"""

    for topkey in objtomerge.keys():
      for subkey in objtomerge[topkey].keys():
        if not topkey in self.tallychars or not subkey in self.tallychars:
           continue
        if subkey in self[topkey].keys():
          self[topkey][subkey]=self[topkey][subkey]+(objtomerge[topkey][subkey]*weight)
        else:
          self[topkey][subkey]=(objtomerge[topkey][subkey]*weight)

  def printtable(self):
    """ Prints the frequency tables as a python function that can be used to lookup the characters that follow a character.  You can plug the resulting script into a python program to lookup the most frequent character to follow another character.  For example you can call "lookupfreq("q")" and it will return a string containing all the characters in frequency order such as "ustrnalq1f"   where "u" is the most frequenct character to follow "q" and "f" is the least frequent.
"""
    print "def lookupfreq(index):"
    print "  return {"
    for keys in sorted(self.keys(),reverse=True, key=self.get):
      print "    \""+str(keys)+"\":",
      letters=""
      for subkeys in sorted(self[keys].keys(),reverse=True, key=self[keys].get):
        letters=letters+subkeys
      print '"'+letters+'",'
    print "  }[index]"

  def lookup(self,letter):
    """ Returns a string of characters in frequency order. """
    letters=""
    for subkeys in sorted(self[letter].keys(),reverse=True, key=self[letter].get):
        letters=letters+subkeys
    return letters


  def probability(self,top,sub):
    """This function will print the probability that a character will follow another.
Example:
counter.probability("q","u") - Will tell you the chance that U will follow Q based on the data in the character frequency counter object.  For example,  If you have a Q there is approximately a 30% chance that the next character is a U."""
    totalcounts=0
    for key in self[top].keys():
      totalcounts += self[top][key]
    print str(totalcounts)+" letters follow the letter "+str(top)+".  Of those "+str(self[top][sub])+" are the letter "+str(sub)+". That is "+str(float(self[top][sub])/float(totalcounts)*100)+" percent."
    return float(self[top][sub])/float(totalcounts)*100

  def printraw(self):
    """Prints the raw python data structure containing the frequency table."""
    print self

  def save(self,filename):
    """Saves the raw python data structure from the file specified.  Save and Load are used to write the data structure to disk so you can come back to it later or exchange them with other developers.
Example:
counter.save("/home/user/savedfreqcounter.txt") - save the data structure from disk.
"""
    file=open(filename,'w')
    file.write(str(self.items()))
    file.close()

  def load(self,filename):
    """Loads the raw python data structure from the file specified. Load is using the raw data structures. Save and Load are used to write the data structure to disk so you can come back to it later or exchange them with other developers.  This is not the same as tallyfile.  Tallyfile analyzes a files character frequencies. 
Example:
counter.load("/home/user/savedfreqcounter.txt") - Loads the data structure from disk.
"""
    file=open(filename,'r')
    self.update(dict(eval(file.read())))
    file.close()

  def promote(self,top,sub,weight):
    """Promotes sub up weight positions inside of tops array. 
Example:
counter.promote("c","a",5) - a will move up 5 places in c's table
"""
    #import pdb;pdb.set_trace()
    chartable=self[top]
    currentoffset=sorted(chartable.keys(),reverse=True, key=chartable.get).index(sub)
    if currentoffset<weight:
      currentoffset=weight
    movebeforeletter=sorted(chartable.keys(),reverse=True, key=chartable.get)[currentoffset-weight]
    addthisamount=self[top][movebeforeletter]-self[top][sub]+1
    self[top][sub]=self[top][sub]+addthisamount    
    return

  def resetcounts(self):
    """Reset the counts for all of the character frequencies to 1 giving every character an equal probability of following other characters."""
    for keys in self.keys():
      for subkeys in self[keys].keys():
        self[keys][subkeys]=1
